from LoadData import load_awa2
import torch
from torch.utils.data import DataLoader
from torch.utils.data.sampler import SubsetRandomSampler
import torchvision
from torchvision import transforms
from torchvision.datasets import ImageFolder
from torchvision.datasets import FashionMNIST, CIFAR10
import numpy as np


def get_awa2_images(batch_size, resize=None, root='Data/AwA2/'):
    train_path = root + "train"
    test_path = root + "test"
    # train_transform= transforms.Compose([
    #     transforms.RandomCrop(224, 8),
    #     transforms.RandomHorizontalFlip(),
    #     transforms.ToTensor(),
    #     transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    # ])
    origin_transform = transforms.Compose([
        transforms.ToTensor(),
    ])
    
    normal_transform = transforms.Compose([
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
    ])
    # print("hello")
    origin_test_dataset = ImageFolder(test_path, transform=origin_transform)
    normal_test_dataset = ImageFolder(test_path, transform=normal_transform)
    sample_size = 128
    np.random.seed(34)
    sample_idx = np.random.choice(range(len(origin_test_dataset)), sample_size)
    # np.random.choice(rorigin_data_loader.size)
    # sampler = SubsetRandomSampler(np.random.choice(range(len(origin_test_dataset)), sample_size))
    sampler = SubsetRandomSampler(sample_idx)
    # print(sampler.indices)
#     sampler.indices = [11681, 5242, 10730, 5993 12022  3157   324  5667  6597  1675  5610  3278
#   7698  1565  4344  6941 12902  8565  2166  2324  4069  3680 13118  1285
#    916  1971  7699   981  6335  2536 11848  4591  8557 12616 13182 10346
#   1101 10404  5413 14562  5306  2850 10601  5178  2750 12454  6872 13634
#   9798  5344  9003 12501  8564  3615 11988  1508  7385 13559 12478  8020
#   2650 12493  8620  4087  2835  2666  9949 12526 11814  8471 13965 14150
#   4232   376   350   889  1124  8740 10090  8925  6320 13837 12529  4097
#  12805  2010  2761  6243 14245  7168 12720 10284  2779  2402   419  3220
#  12945  1350  1131 10727  4387  2982  8341  4241  8765  4107 13830  6203
#   9805  3741  9151 11025 10207  5120 10325 12902  7271 10283 14508  4318
#    945  1190  5143  7010  3962 10729  5221  7282]
#     exit(0)
    origin_data_loader = DataLoader(origin_test_dataset, batch_size=batch_size, sampler=sampler, shuffle=False, num_workers=4)
    normal_data_loader = DataLoader(normal_test_dataset, batch_size=batch_size, sampler=sampler, shuffle=False, num_workers=4)
    # print("hello")
    
    origin_images = None
    normal_images = None
    labels = None
    for i, ((img1, _), (img2, lab)) in enumerate(zip(origin_data_loader, normal_data_loader)):
        origin_images = img1
        normal_images = img2
        labels = lab
        break
    origin_images = origin_images.numpy().transpose((0, 2, 3, 1))
    return origin_images, normal_images, labels


if __name__ == '__main__':
    get_awa2_images(128, 224)